package com.bxmdm.ragdemo.reader;

import com.fasterxml.jackson.core.type.TypeReference;
import com.fasterxml.jackson.databind.ObjectMapper;
import org.springframework.ai.document.Document;
import org.springframework.ai.document.DocumentReader;
import org.springframework.ai.embedding.EmbeddingClient;
import org.springframework.core.io.Resource;
import org.springframework.util.StreamUtils;

import java.io.IOException;
import java.nio.charset.StandardCharsets;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.UUID;

/**
 * @Description: 简单英语问答数据读取器（非加权版本），用于对比测试
 */
public class SimpleEnglishQAReader implements DocumentReader {

    private final Resource resource;
    private final ObjectMapper objectMapper = new ObjectMapper();
    private final EmbeddingClient embeddingClient;

    // 文档缓存
    private List<Document> documentCache = null;

    public SimpleEnglishQAReader(Resource resource, EmbeddingClient embeddingClient) {
        this.resource = resource;
        this.embeddingClient = embeddingClient;
    }

    @Override
    public List<Document> get() {
        // 检查缓存
        if (documentCache != null) {
            return documentCache;
        }

        try {
            List<Document> documents = new ArrayList<>();
            String content = StreamUtils.copyToString(this.resource.getInputStream(), StandardCharsets.UTF_8);

            // 从JSON中解析问答列表
            List<Map<String, Object>> qaItems = objectMapper.readValue(
                    content,
                    new TypeReference<List<Map<String, Object>>>() {
                    });

            for (Map<String, Object> qaItem : qaItems) {
                Document doc = processQAPair(qaItem);
                if (doc != null) {
                    documents.add(doc);
                }
            }

            // 缓存结果
            documentCache = documents;
            return documents;
        } catch (IOException e) {
            throw new RuntimeException("读取问答数据文件失败: " + e.getMessage(), e);
        }
    }

    @SuppressWarnings("unchecked")
    private Document processQAPair(Map<String, Object> qaItem) {
        // 提取基本信息
        String postId = qaItem.getOrDefault("post_id", "").toString();
        String question = qaItem.getOrDefault("question", "").toString();
        String answer = qaItem.getOrDefault("answer", "").toString();
        Integer score = qaItem.containsKey("score") ? Integer.valueOf(qaItem.get("score").toString()) : 0;
        List<String> tags = qaItem.containsKey("tags") ? (List<String>) qaItem.get("tags") : new ArrayList<>();
        Integer answerCount = qaItem.containsKey("answer_count")
                ? Integer.valueOf(qaItem.get("answer_count").toString())
                : 0;
        List<String> allAnswers = qaItem.containsKey("all_answers") ? (List<String>) qaItem.get("all_answers")
                : new ArrayList<>();

        // 如果问题为空，跳过
        if (question.isEmpty()) {
            return null;
        }

        // 使用简单的非加权向量计算
        List<Double> embedding = null;
        try {
            // 直接对原始问题进行向量化，不进行任何加权处理
            String cleanQuestion = question.replaceAll("\\s+", " ").trim();
            embedding = embeddingClient.embed(cleanQuestion);
        } catch (Exception e) {
            System.err.println("计算问题嵌入向量失败: " + e.getMessage());
            return null;
        }

        // 构建元数据
        Map<String, Object> metadata = new HashMap<>();
        metadata.put("post_id", postId);
        metadata.put("score", score);
        metadata.put("question", question);
        metadata.put("answer", answer);
        metadata.put("tags", tags);
        metadata.put("answer_count", answerCount);
        metadata.put("all_answers", allAnswers);
        metadata.put("embedding", embedding);

        // 创建Document对象 - 使用原始问题内容
        String id = UUID.randomUUID().toString();
        String cleanQuestion = question.replaceAll("\\s+", " ").trim();
        Document document = new Document(id, cleanQuestion, metadata);
        return document;
    }
}
